from datetime import datetime
import os
from bs4 import BeautifulSoup
from web.models import Leak


def load(file_dir):
    """数据加载 时间大概在10min 信息在16w左右"""
    file_list = os.listdir(file_dir)
    for file in file_list:
        with open(file_dir + "/" + file, "r", encoding='utf-8') as f:
            lines = f.readlines()
            data = "".join(lines)
        print(f"== {file}加载成功！")
        soup = BeautifulSoup(data, "lxml")
        for item in soup.select("entry"):
            v_id = item.select("vuln-id")[0].text.strip()
            try:
                Leak.objects.create(
                    name=item.select("name")[0].text.strip(),
                    v_id=v_id,
                    published=datetime.strptime(item.select("published")[0].text.strip(), "%Y-%m-%d"),
                    modified=datetime.strptime(item.select("modified")[0].text.strip(), "%Y-%m-%d"),
                    severity=item.select("severity")[0].text.strip(),
                    v_type=item.select("vuln-type")[0].text.strip(),
                    desc=item.select("vuln-descript")[0].text.strip(),
                )
            except ValueError:
                continue
            print(f"{v_id}信息数据库写入成功！")


if __name__ == '__main__':
    load(r"D:\Python_Project\zp\web\static\data")
